#include "arp.h"

#include "ethernet.h"
#include "net.h"

#include <stdio.h>
#include <string.h>
/**
 * @brief 初始的arp包
 *
 */
static const arp_pkt_t arp_init_pkt = {
    .hw_type16 = swap16(ARP_HW_ETHER),
    .pro_type16 = swap16(NET_PROTOCOL_IP),
    .hw_len = NET_MAC_LEN,
    .pro_len = NET_IP_LEN,
    .sender_ip = NET_IF_IP,
    .sender_mac = NET_IF_MAC,
    .target_mac = {0}};

/**
 * @brief arp 地址转换表，<ip,mac>
 *
 */
map_t arp_table;

/**
 * @brief arp 缓冲区，<ip,buf_t>
 *
 * 暂存 “不知道目标 MAC 地址而暂时无法发送” 的 IP 数据包。
 */
map_t arp_buf;

/**
 * @brief 打印一条arp表项
 *
 * @param ip 表项的ip地址
 * @param mac 表项的mac地址
 * @param timestamp 表项的更新时间
 */
void arp_entry_print(void *ip, void *mac, time_t *timestamp) {
    printf("%s | %s | %s\n", iptos(ip), mactos(mac), timetos(*timestamp));
}

/**
 * @brief 打印整个arp表
 *
 */
void arp_print() {
    printf("===ARP TABLE BEGIN===\n");
    map_foreach(&arp_table, arp_entry_print);
    printf("===ARP TABLE  END ===\n");
}

/**
 * @brief 发送一个arp请求
 *
 * @param target_ip 想要知道的目标的ip地址
 */
void arp_req(uint8_t *target_ip) {
    // TO-DO
    // Step1. 初始化缓冲区，对 txbuf 进行初始化。
    if (buf_init(&txbuf, sizeof(arp_pkt_t)) == -1)
        return;  // 失败

    // Step2. 填写ARP报头
    arp_pkt_t *header = (arp_pkt_t *)txbuf.data;
    *header = arp_init_pkt;
    header->opcode16 = swap16(ARP_REQUEST);            // 报文类型：请求/响应
    memcpy(header->target_ip, target_ip, NET_IP_LEN);  // 目的IP

    // Step3. 发送 ARP 请求（广播）
    // 目标 MAC 地址设置为广播地址：FF - FF - FF - FF - FF - FF。
    ethernet_out(&txbuf, ether_broadcast_mac, NET_PROTOCOL_ARP);
}

/**
 * @brief 发送一个arp响应
 *
 * @param target_ip 目标ip地址
 * @param target_mac 目标mac地址
 */
void arp_resp(uint8_t *target_ip, uint8_t *target_mac) {
    // TO-DO
    // Step1. 初始化缓冲区，对 txbuf 进行初始化。
    if (buf_init(&txbuf, sizeof(arp_pkt_t)) == -1)
        return;  // 失败

    // Step2. 填写 ARP 报头首部
    arp_pkt_t *header = (arp_pkt_t *)txbuf.data;
    header->opcode16 = swap16(ARP_REPLY);                 // 报文类型
    memcpy(header->target_ip, target_ip, NET_IP_LEN);     // 目的IP
    memcpy(header->target_mac, target_mac, NET_MAC_LEN);  // 目的IP

    // Step3. 发送 ARP 报文（单播）
    ethernet_out(&txbuf, target_mac, NET_PROTOCOL_ARP);
}

/**
 * @brief 处理一个收到的数据包
 *
 * @param buf 要处理的数据包
 * @param src_mac 源mac地址
 */
void arp_in(buf_t *buf, uint8_t *src_mac) {
    // TO-DO
    // Step1. 检查数据长度
    if (buf->len < sizeof(arp_pkt_t)) {
        return;
    }

    // Step2. 查看报文是否完整
    arp_pkt_t *data = (arp_pkt_t *)buf->data;
    if (swap16(data->hw_type16) != ARP_HW_ETHER ||
        swap16(data->pro_type16) != NET_PROTOCOL_IP ||
        data->hw_len != NET_MAC_LEN ||
        data->pro_len != NET_IP_LEN) {
        return;  // 非法 ARP 报文，丢弃
    }

    // Step3. 调用 map_set() 函数更新 arp_table
    map_set(&arp_table, data->sender_ip, data->sender_mac);

    // Step4. 查看该报文的 IP 地址是否有对应的 arp_buf 缓存。
    if (map_get(&arp_buf, data->sender_ip) != NULL) {
        // 若有缓存，说明 ARP 分组队列里面有待发送的数据包。
        // 即调用 arp_out() 发送来自 IP 层的数据包时，由于没有找到对应的 MAC 地址而先发送了 ARP request 报文，此时收到了该 request 的应答报文。

        // 将缓存的数据包 arp_buf.get(data->sender_ip) 发送给以太网层，
        // 调用 map_delete() 函数将arp_buf内对应项删除。
        buf_t *buf_to_send = map_get(&arp_buf, data->sender_ip);
        ethernet_out(buf_to_send, data->sender_mac, NET_PROTOCOL_IP);
        map_delete(&arp_buf, data->sender_ip);
    } else {
        // 若该接收报文的 IP 地址没有对应的 arp_buf 缓存，
        // 判断是否为 ARP_REQUEST 请求，且 target_ip 是本机的 IP。
        // 若是，则认为是请求本主机 MAC 地址的 ARP 请求报文，回应一个响应报文。
        if (swap16(data->opcode16) == ARP_REQUEST && memcmp(data->target_ip, net_if_ip, NET_IP_LEN) == 0) {
            arp_resp(data->sender_ip, data->sender_mac);
        }
    }
}

/**
 * @brief 处理一个要发送的数据包
 *
 * @param buf 要处理的数据包
 * @param ip 目标ip地址
 * @param protocol 上层协议
 */
void arp_out(buf_t *buf, uint8_t *ip) {
    // TO-DO
    // Step1. 查找 ARP 表
    void *mac = map_get(&arp_table, ip);

    // Step2. 若能找到该 IP 地址对应的 MAC 地址，则将数据包直接发送给以太网层
    if (mac != NULL) {
        ethernet_out(buf, mac, NET_PROTOCOL_IP);
    }

    // Step3. 若未找到对应的 MAC 地址，
    else {
        // 判断 arp_buf 中是否已经有包：
        // 若有包，说明正在等待该 IP 的回应 ARP 请求，不能再发送 ARP 请求；

        // 若没有包，将来自 IP 层的数据包缓存到 arp_buf 中，
        // 然后发送ARP request 报文：请求目标 IP 地址对应的 MAC 地址
        if (map_get(&arp_buf, ip) == NULL) {
            buf_t buf_cpy;
            buf_copy(&buf_cpy, buf, 0);
            map_set(&arp_buf, ip, &buf_cpy);
            arp_req(ip);
        }
    }

    // 扩展功能（未实现）
    // 由于 map 数据结构存在一定限制，相同 IP 的 arp_buf 长度仍为 1，同学们可自行设计数据结构以实现不丢包的扩展功能。
}

/**
 * @brief 初始化arp协议
 *
 */
void arp_init() {
    // 初始化
    // arp 地址转换表，arp_table <ip,mac>
    // arp 缓冲区，arp_buf <ip,buf_t>
    map_init(&arp_table, NET_IP_LEN, NET_MAC_LEN, 0, ARP_TIMEOUT_SEC, NULL, NULL);
    map_init(&arp_buf, NET_IP_LEN, sizeof(buf_t), 0, ARP_MIN_INTERVAL, NULL, buf_copy);

    // 将 <ARP协议,处理程序> 注册到net_table中
    net_add_protocol(NET_PROTOCOL_ARP, arp_in);

    // 广播发送自己的 ARP 包
    arp_req(net_if_ip);
}